import copy
import time

import torch
from torchvision.datasets import FashionMNIST
from torchvision import transforms
import torch.utils.data as Data
import numpy as np
import matplotlib.pyplot as plt
from model import LeNet#导入我们建的模型
import torch.nn as nn
import pandas as pd

#数据加载函数
def train_val_data_process():#处理训练集和验证集
    train_data = FashionMNIST('./data',
                              train=True,
                              transform=transforms.Compose([transforms.Resize(size=28), transforms.ToTensor()]),
                              download=True)
    train_data,val_data=Data.random_split(train_data,[round(0.8*len(train_data)),round(0.2*len(train_data))])

    train_dataloader = Data.DataLoader(dataset=train_data,
                                   batch_size=32,
                                   shuffle=True,
                                   num_workers=4)

    val_dataloader = Data.DataLoader(dataset=val_data,
                                   batch_size=32,
                                   shuffle=True,
                                   num_workers=4)
    return train_dataloader,val_dataloader

#模型训练函数
def train_model_process(model,train_dataloader,val_dataloader,num_epochs):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)#优化器
    criterion = nn.CrossEntropyLoss()#分类一般用交叉熵损失函数

    model=model.to(device)#模型放到设备里
    #这一行代码会有显著的影响，因为这行代码的作用是在训练开始时保存模型的初始权重
    best_model_wts = copy.deepcopy(model.state_dict())

    #初始化参数
    best_acc = 0.0#最高精确度
    train_loss_all=[]#训练集损失函数列表
    val_loss_all=[]#验证集损失函数列表

    train_acc_all=[]#训练集准确度列表
    val_acc_all=[]#验证集准确度列表

    since=time.time()#当前时间保存下来

    for epoch in range(num_epochs):#一轮有很多批次
        print('Epoch {}/{}'.format(epoch, num_epochs-1))#打印训练轮次
        print('-' * 10)#虚线

        #初始化参数
        train_loss=0.0#训练集损失函数
        train_corrects=0.0#训练集准确度
        val_loss=0.0#验证集损失函数
        val_corrects=0.0#验证集准确度

        train_num=0#训练集样本数量
        val_num=0

        #train_dataloader第一次循环是第一个批次的数据；第二次循环是第二个批次的数据
        #batchsize是32，所以b_x是32*28*28*1；b_y是32*label
        for step,(b_x,b_y) in enumerate(train_dataloader):
            b_x=b_x.to(device)#特征放到设备里
            b_y=b_y.to(device)#label放到设备里
            #训练中，模型和数据都要放到设备里。

            model.train()#模型打开训练模式

            #前向传播过程，输入为一个batch，输出为一个batch中对应的预测
            output=model(b_x)#前向传播

            pre_lab=torch.argmax(output,dim=1)#查找每一行中最大值对应的行标
            loss=criterion(output,b_y)#计算每一个batch的损失

            optimizer.zero_grad()#将梯度初始化为0
            loss.backward()#反向传播计算
            optimizer.step()#根据网络反向传播的梯度信息来更新网络的参数，以起到降低loss函数计算值的作用

            #这里的loss是每个批次中每个样本的平均值。train_loss是所有训练样本的损失
            train_loss += loss.item()*b_x.size(0)#对损失函数进行累加，b_x.size(0) 给出了当前batch的样本数量
            train_corrects += torch.sum(pre_lab==b_y.data)#train_corrects是该轮次中所有预测正确分数目

            train_num += b_x.size(0)

        #验证过程
        for step,(b_x,b_y) in enumerate(val_dataloader):
            b_x=b_x.to(device)#特征放到设备里
            b_y=b_y.to(device)#label放到设备里

            model.eval()#设置模型为评估模式

            #前向传播过程，输入为一个batch，输出为一个batch中对应的预测
            output=model(b_x)#前向传播,输出(batch_size, 10)

            pre_lab=torch.argmax(output,dim=1)#查找每一行中最大值对应的行标
            loss=criterion(output,b_y)#计算每一个batch的损失

            # 这里的loss是每个批次中每个样本的平均值。val_loss是所有训练样本的损失
            val_loss += loss.item()*b_x.size(0)#对损失函数进行累加，b_x.size(0) 给出了当前batch的样本数量
            val_corrects += torch.sum(pre_lab==b_y.data)#如果预测正确，则准确度train_corrects加1

            val_num += b_x.size(0)

        #计算并保存每一次迭代的loss值和准确率
        train_loss_all.append(train_loss/train_num)#每个轮次平均的loss值
        #计算并保存每轮次训练集的准确率
        train_acc_all.append(train_corrects.double().item()/train_num)#double取双精度
        val_loss_all.append(val_loss/val_num)
        val_acc_all.append(val_corrects.double().item() / val_num)

        #-1表示列表中最后一个
        print('{}Train Loss: {:.4f} Train Acc: {:.4f}'.format(epoch, train_loss_all[-1], train_acc_all[-1]))
        print('{}Val Loss: {:.4f} Val Acc: {:.4f}'.format(epoch, val_loss_all[-1], val_acc_all[-1]))

        #寻找最高准确度的权重参数
        if val_acc_all[-1]>best_acc:
            best_acc=val_acc_all[-1]
            best_model_wts = copy.deepcopy(model.state_dict())
        #训练耗费时间
        time_used=time.time()-since
        print("训练和验证耗费的时间{:.0f}m{:.0f}s".format(time_used//60, time_used%60))

    #选择最优参数
    #加载最高准确率下的模型参数
    torch.save(best_model_wts,r'D:\My_File\Gitee\deep-learning\pytorch框架与经典卷积神经网络_炮哥带你学\LeNet\best_model.pth')#保存参数pth格式，调用更快，使用时只用调用模型框架就行

    train_process=pd.DataFrame(data={'epoch':range(num_epochs),
                                     'train_loss_all':train_loss_all,
                                     'val_loss_all':val_loss_all,
                                     'train_acc_all':train_acc_all,
                                     'val_acc_all':val_acc_all
                                     })
    return train_process

def matplot_acc_loss(train_process):
    plt.figure(figsize=(12,4))
    plt.subplot(1, 2, 1)#1行2列第一张图
    plt.plot(train_process['epoch'],train_process.train_loss_all,'ro-',label='train loss')
    plt.plot(train_process['epoch'],train_process.val_loss_all,'bs-',label='val loss')
    plt.legend()
    plt.xlabel('epoch')
    plt.ylabel('loss')

    plt.figure(figsize=(12, 4))
    plt.subplot(1, 2, 2)
    plt.plot(train_process['epoch'], train_process.train_acc_all, 'ro-', label='train acc')
    plt.plot(train_process['epoch'], train_process.val_acc_all, 'bs-', label='val acc')
    plt.legend()
    plt.xlabel('epoch')
    plt.ylabel('acc')
    plt.show()

if __name__ == '__main__':
    #将模型实例化
    LeNet=LeNet()#实例化一个模型（如 LeNet = LeNet()），PyTorch 会自动调用模型中的 __init__ 方法，并初始化模型的参数
    train_dataloader,val_dataloader=train_val_data_process()
    train_process=train_model_process(LeNet,train_dataloader,val_dataloader,20)#训练
    matplot_acc_loss(train_process)
